#! /usr/bin/env python
#coding=utf-8

import os
import json

class ImageExtractor(object):
	def __init__(self, img_path, output_dir):
		self._img_name = None
		self._img_file = None
		self._img_path = img_path
		file_name = self.__get_img_file(img_path)
		if not file_name:
			return
		img_name = os.path.splitext(file_name)[0]
		if img_name not in ("system", "vendor", "chipset", "chip_prod", "sys_prod"):
			return

		self._img_name = img_name
		self._img_file = os.path.join(img_path, file_name)
		self._output_dir = output_dir

	def __get_img_file(self, img_path):
		files = os.listdir(img_path)

		for f in files:
			if f.endswith(".img"):
				return f

		return None

	def get_image_name(self):
		return self._img_name

	def get_image_path(self):
		return self._img_path

	def get_extracted_image_path(self):
		top_dir = os.path.join(self._output_dir, "packages/phone", self._img_name, self._img_name)
		if os.path.exists(os.path.join(top_dir, self._img_name)):
			return os.path.join(top_dir, self._img_name)
		return top_dir

	def extract_image(self):
		if not self._img_name:
			return

		mount_dir = os.path.join(self._output_dir, "packages/phone/mount", self._img_name)

		# Create directory
		try:
			os.makedirs(mount_dir)
		except:
			pass

		try:
			os.system("sudo umount %s" % mount_dir)
		except:
			pass

		# Do mount
		print("sudo mount -t ext4 -o loop %s %s" % (self._img_file, mount_dir))
		os.system("sudo mount -t ext4 -o loop %s %s" % (self._img_file, mount_dir))

		# chown
		try:
			curUser = os.getlogin()
			os.system("sudo chown -R %s:%s %s" % (curUser, curUser, mount_dir))
		except:
			pass

		# Do soft link
		sub_dir = os.path.join(mount_dir, self._img_name)
		if os.path.exists(sub_dir):
			os.system("ln -sf mount/%s/%s %s" % (self._img_name, self._img_name, os.path.join(self._output_dir, "packages/phone", self._img_name)))
		else:
			os.system("ln -sf mount/%s %s" % (self._img_name, os.path.join(self._output_dir, "packages/phone", self._img_name)))

	def cleanup(self):
		if not self._img_name:
			return

		mount_dir = os.path.join(self._output_dir, "packages/phone/mount", self._img_name)
		print("sudo umount %s" % mount_dir)
		os.system("sudo umount %s" % mount_dir)

		os.system("sudo rm -fr %s" % mount_dir)

class SystemImageProcessor(object):
	def __init__(self, args):
		if not args.extract:
			self._extract = False
			args.input = args.input[0]
			return

		if not args.output:
			print("Output directory is required")
			sys.exit()

		self._extract = True
		self._output_dir = args.output

		self._imgs = []
		for i in args.input:
			self._imgs.append(ImageExtractor(i, args.output))

		args.input = args.output

	def __merge_module_info(self):
		infos = []
		for img in self._imgs:
			img_name = img.get_image_name()
			print("Load %s module info now ..." % img_name)
			try:
				with open(os.path.join(img.get_image_path(), "system_module_info.json"), "r") as f:
					info = json.load(f)
			except:
				continue
			for item in info:
				for dest in item["dest"]:
					if dest.startswith(img_name):
						infos.append(item)

		file_dir = os.path.join(self._output_dir, "packages/phone")
		try:
			os.makedirs(file_dir)
		except:
			pass

		with open(os.path.join(file_dir, "system_module_info.json"), "w") as f:
			f.write(json.dumps(infos, indent=4))

	def __merge_inner_api(self):
		infos = None
		for img in self._imgs:
			img_name = img.get_image_name()
			print("Load %s innerapi info now ..." % img_name)
			try:
				with open(os.path.join(img.get_image_path(), "inner_kits_info.json"), "r") as f:
					info = json.load(f)
			except:
				continue
			if not infos:
				infos = info
				continue

			for component, innerapis in info.items():
				if component not in infos:
					infos[component] = innerapis
					continue
				for name, innerapi in innerapis.items():
					if name not in infos[component]:
						infos[component][name] = innerapi

		file_dir = os.path.join(self._output_dir, "build_configs/parts_info")
		try:
			os.makedirs(file_dir)
		except:
			pass

		with open(os.path.join(file_dir, "inner_kits_info.json"), "w") as f:
			f.write(json.dumps(infos, indent=4))

	def pre_process(self):
		if not self._extract:
			return

		# Extract image files
		for img in self._imgs:
			img.extract_image()

		# Merge system_module_info.json
		self.__merge_module_info()

		# Merge inner_kits_info.json
		self.__merge_inner_api()

	def cleanup(self):
		if not self._extract:
			return

		# Extract image files
		for img in self._imgs:
			img.cleanup()

if __name__ == '__main__':
	import argparse

	parser = argparse.ArgumentParser(description='Collect architecture information from asset files.')
	parser.add_argument('-i', '--input', help='input asset files root directory', action='append', required=True)
	parser.add_argument('-o', '--output', help='output architecture information database directory', required=False)
	parser.add_argument('-x', '--extract', help='need extract image file system', required=False, default=False, action='store_true')

	args = parser.parse_args()
	pro = SystemImageProcessor(args)
	pro.pre_process()
	pro.cleanup()
